import subprocess
import asyncio
from neutron import Plugin, OS, User
import os
from typing import Literal
from fastapi import WebSocket
from starlette import websockets
import uuid
from . import icon

SEARCH_MAX_LEN = 2000
REGEX_MODE: Literal["--regexp", "--regex"] = "--regex"
SEARCH_ENCODE = "utf-8"


async def timeout_second(tasks):  # 防超时
    try:
        return await asyncio.wait_for(asyncio.shield(asyncio.create_task(asyncio.to_thread(tasks))), timeout=1)
    except TimeoutError:
        raise TimeoutError("任务超时")


def get_size(path):  # 获取文件大小并添加单位
    if (os.path.isdir(path)):
        return ""
    try:
        size_bytes = os.stat(path).st_size
        if (size_bytes > 1024*1024*1024*1024):
            return f"{size_bytes/(1024*1024*1024*1024):.2f}TiB"
        elif (size_bytes > 1024*1024*1024):
            return f"{size_bytes/(1024*1024*1024):.2f}GiB"
        elif (size_bytes > 1024*1024):
            return f"{size_bytes/(1024*1024):.2f}MiB"
        elif (size_bytes > 1024):
            return f"{size_bytes/(1024):.2f}KiB"
        else:
            return f"{size_bytes}B"
    except:
        return ""


async def create_db(root_dir, db_dir):  # 构建数据库
    try:
        run = subprocess.Popen(
            ["updatedb", '-o', os.path.abspath(db_dir), "-l", "no", '-U', root_dir], shell=False, stdout=None, stderr=None)
    except FileNotFoundError:
        return 127
    try:
        ret = await asyncio.to_thread(lambda: run.wait())
    except asyncio.CancelledError:
        run.kill()
        raise
    return ret


async def search_db(db_dir, root_dir, file_name, reg_mode=False):

    def sort_key(x):  # 排序
        # 文件夹优先
        return ("0" if os.path.isdir(x["fullpath"]) else "1")+x["name"]

    try:
        run = subprocess.Popen(
            ["plocate", "-i", "-b", file_name, "-l", str(SEARCH_MAX_LEN), "-d", db_dir]+([REGEX_MODE] if reg_mode else []), shell=False, stdout=subprocess.PIPE, stderr=None)
    except FileNotFoundError:
        return ([], 127)
    try:  # 进行搜索
        outlist = (await asyncio.to_thread(lambda: run.communicate()))[0].decode(SEARCH_ENCODE).splitlines()
    except asyncio.CancelledError:
        run.kill()
        raise
    outlist = [
        {
            "img": "icons/"+(icon.get_icon(i)),
            "name": (i[len(root_dir)+1:].lstrip("/") if i.startswith(root_dir+os.sep) else i),
            "size": await timeout_second(lambda: get_size(i)),
            "isdir": await timeout_second(lambda: os.path.isdir(i)),
            "fullpath": i
        } for i in outlist
    ]
    outlist = await timeout_second(lambda: sorted(outlist, key=sort_key))
    return (outlist, run.returncode)


async def callback_sys(coro, callback):
    ret = await coro
    await (callback(ret))
    return ret


async def search_ws(user: User, os: OS, app: Plugin,  websocket: WebSocket, search_path: str):
    temp = await app.filestorage.load(user, "databases", create_everytime=True)
    db_path = temp+uuid.uuid1().hex+".db"  # 临时数据库
    initdb_task = None
    search_task = None
    closed = False

    try:
        await websocket.accept()

        search_str = ""

        async def search_callback(ret):  # 搜索回调
            nonlocal closed
            if (closed):
                return
            await websocket.send_json({
                "type": "result",
                "data": {
                    "retcode": ret[1],
                    "data": ret[0]
                }
            })

        async def task_callback(ret):  # 初始化数据库回调
            nonlocal search_str, search_task, closed

            if (closed):
                return

            await websocket.send_json({
                "type": "db_created",
                "data": {
                    "data": ret
                }
            })
            if (search_str != ""):
                if (search_task is not None and not search_task.done()):
                    search_task.cancel()
                search_task = asyncio.create_task(
                    callback_sys(search_db(db_path, search_path, search_str), search_callback))
        initdb_task = asyncio.create_task(callback_sys(
            create_db(search_path, db_path), task_callback))

        search_task = asyncio.create_task(create_db(search_path, db_path))

        while True:
            data = await websocket.receive_json()
            if (data["type"] == "search"):
                if (not initdb_task.done()):
                    search_str = data["data"]
                else:
                    search_task = asyncio.create_task(
                        callback_sys(search_db(db_path, search_path, data["data"]), search_callback))
    except (RuntimeError, websockets.WebSocketDisconnect):
        if (initdb_task is not None and not initdb_task.done()):
            initdb_task.cancel()
        if (search_task is not None and not search_task.done()):
            search_task.cancel()
        if (await asyncio.to_thread(lambda: os.path.exists(db_path))):
            await asyncio.to_thread(lambda: os.remove(db_path))
        closed = True  # 设置连接关闭标志
